
/**
 * @file exponent.c
 * A/52 exponent functions
 */

#include "common.h"

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <assert.h>

#include "exponent.h"
#include "a52.h"

/**
 * LUT for number of exponent groups present.
 * expsizetab[exponent strategy][number of coefficients]
 */
static uint8_t nexpgrptab[2][3][256] = {{{0}}};

/**
 * Pre-defined sets of exponent strategies. A strategy set is selected for
 * each channel in a frame.
 */
static const uint8_t str_predef[A52_EXPSTR_SETS][6] = {
    { EXP_D15, EXP_REUSE, EXP_REUSE, EXP_REUSE, EXP_REUSE, EXP_REUSE },
    { EXP_D15, EXP_REUSE, EXP_REUSE, EXP_REUSE, EXP_REUSE,   EXP_D45 },
    { EXP_D15, EXP_REUSE, EXP_REUSE, EXP_REUSE,   EXP_D25, EXP_REUSE },
    { EXP_D15, EXP_REUSE, EXP_REUSE, EXP_REUSE,   EXP_D45,   EXP_D45 },
    { EXP_D25, EXP_REUSE, EXP_REUSE,   EXP_D25, EXP_REUSE, EXP_REUSE },
    { EXP_D25, EXP_REUSE, EXP_REUSE,   EXP_D25, EXP_REUSE,   EXP_D45 },
    { EXP_D25, EXP_REUSE, EXP_REUSE,   EXP_D45,   EXP_D25, EXP_REUSE },
    { EXP_D25, EXP_REUSE, EXP_REUSE,   EXP_D45,   EXP_D45,   EXP_D45 },
    { EXP_D25, EXP_REUSE,   EXP_D15, EXP_REUSE, EXP_REUSE, EXP_REUSE },
    { EXP_D25, EXP_REUSE,   EXP_D25, EXP_REUSE, EXP_REUSE,   EXP_D45 },
    { EXP_D25, EXP_REUSE,   EXP_D25, EXP_REUSE,   EXP_D25, EXP_REUSE },
    { EXP_D25, EXP_REUSE,   EXP_D25, EXP_REUSE,   EXP_D45,   EXP_D45 },
    { EXP_D25, EXP_REUSE,   EXP_D45,   EXP_D25, EXP_REUSE, EXP_REUSE },
    { EXP_D25, EXP_REUSE,   EXP_D45,   EXP_D25, EXP_REUSE,   EXP_D45 },
    { EXP_D25, EXP_REUSE,   EXP_D45,   EXP_D45,   EXP_D25, EXP_REUSE },
    { EXP_D25, EXP_REUSE,   EXP_D45,   EXP_D45,   EXP_D45,   EXP_D45 },
    { EXP_D45,   EXP_D15, EXP_REUSE, EXP_REUSE, EXP_REUSE, EXP_REUSE },
    { EXP_D45,   EXP_D15, EXP_REUSE, EXP_REUSE, EXP_REUSE,   EXP_D45 },
    { EXP_D45,   EXP_D25, EXP_REUSE, EXP_REUSE,   EXP_D25, EXP_REUSE },
    { EXP_D45,   EXP_D25, EXP_REUSE, EXP_REUSE,   EXP_D45,   EXP_D45 },
    { EXP_D45,   EXP_D25, EXP_REUSE,   EXP_D25, EXP_REUSE, EXP_REUSE },
    { EXP_D45,   EXP_D25, EXP_REUSE,   EXP_D25, EXP_REUSE,   EXP_D45 },
    { EXP_D45,   EXP_D25, EXP_REUSE,   EXP_D45,   EXP_D25, EXP_REUSE },
    { EXP_D45,   EXP_D25, EXP_REUSE,   EXP_D45,   EXP_D45,   EXP_D45 },
    { EXP_D45,   EXP_D45,   EXP_D15, EXP_REUSE, EXP_REUSE, EXP_REUSE },
    { EXP_D45,   EXP_D45,   EXP_D25, EXP_REUSE, EXP_REUSE,   EXP_D45 },
    { EXP_D45,   EXP_D45,   EXP_D25, EXP_REUSE,   EXP_D25, EXP_REUSE },
    { EXP_D45,   EXP_D45,   EXP_D25, EXP_REUSE,   EXP_D45,   EXP_D45 },
    { EXP_D45,   EXP_D45,   EXP_D45,   EXP_D25, EXP_REUSE, EXP_REUSE },
    { EXP_D45,   EXP_D45,   EXP_D45,   EXP_D25, EXP_REUSE,   EXP_D45 },
    { EXP_D45,   EXP_D45,   EXP_D45,   EXP_D45,   EXP_D25, EXP_REUSE },
    { EXP_D45,   EXP_D45,   EXP_D45,   EXP_D45,   EXP_D45,   EXP_D45 }
};

/**
 * Pre-defined strategy set indices, sorted most to least common.
 */
static const uint8_t str_predef_priority[A52_EXPSTR_SETS] = {
     2,  8, 10, 17, 24,  3, 14, 13, 22, 21, 26, 11, 30, 15, 23, 31,
    29,  0, 27, 16,  1,  6, 12, 28, 20,  7,  5,  9, 18, 25, 19,  4
};

uint16_t a52_exp_strategy_set_bits_tab[A52_EXPSTR_SETS][256] = {{0}};

/**
 * Initialize exponent tables
 */
void a52_exponent_init(void)
{
    int cpl, i, j, grpsize, ngrps, nc, blk;

    /* init table for number of exponent groups */
    for (cpl = 0; cpl <= 1; cpl++) {
        for (i = 1; i < 4; i++) {
            for (j = 0; j < 256; j++) {
                grpsize = i + (i == EXP_D45);
                ngrps = 0;
                if (cpl) {
                    ngrps = j / (grpsize * 3);
                } else {
                    if (j == 7)
                        ngrps = 2;
                    else
                        ngrps = (j + (grpsize * 3) - 4) / (3 * grpsize);
                }
                nexpgrptab[cpl][i-1][j] = ngrps;
            }
        }
    }

    /* init table for counting exponent bits */
    for (i = 0; i < A52_EXPSTR_SETS; i++) {
        uint16_t *expbits = a52_exp_strategy_set_bits_tab[i];
        for (nc = 0; nc <= 253; nc++) {
            uint16_t bits = 0;
            for (blk = 0; blk < A52_NUM_BLOCKS; blk++) {
                uint8_t es = str_predef[i][blk];
                if (es != EXP_REUSE)
                    bits += (4 + (nexpgrptab[0][es-1][nc] * 7));
            }
            expbits[nc] = bits;
        }
    }
}

/**
 * Get minimum of 2 arrays.
 * Sets exp[i] to min(exp[i], exp1[i])
 */
static void exponent_min(uint8_t *exp, uint8_t *exp1, int n)
{
    int i;
    for (i = 0; i < n; i++)
        exp[i] = MIN(exp[i], exp1[i]);
}

/**
 * Update the exponents so that they are the ones the decoder will decode.
 * Constrain DC exponent, group exponents based on strategy, constrain delta
 * between adjacent exponents to +2/-2.
 */
static void encode_exp_blk_ch(uint8_t *exp, int start, int end,
                              int exp_strategy, int cpl)
{
    int grpsize, ngrps, i, k, exp_min1, exp_min2, abs=0;
    uint8_t exp1[256];
    uint8_t v;

    assert(exp_strategy > EXP_REUSE && exp_strategy <= EXP_D45);

    ngrps = nexpgrptab[cpl][exp_strategy-1][end-start] * 3;
    grpsize = exp_strategy + (exp_strategy == EXP_D45);
    if (cpl)
        abs = start - 1;

    // for D15 strategy, there is no need to group/ungroup exponents
    if (grpsize == 1) {
        // constraint for DC exponent
        if (!cpl)
            exp[abs] = MIN(exp[abs], 15);

        // Decrease the delta between each groups to within 2
        // so that they can be differentially encoded
        for (i = 1; i <= ngrps-cpl; i++)
            exp[start+i] = MIN(exp[start+i], exp[start+i-1]+2);
        for (i = ngrps-1-cpl; i >= 0; i--)
            exp[start+i] = MIN(exp[start+i], exp[start+i+1]+2);

        // coupling absolute exponent
        if (cpl)
            exp[abs] = exp[start] - (exp[start] & 1);

        return;
    }

    // for each group, compute the minimum exponent
    if (grpsize == 2) {
        for (i = 0, k = abs+1; i < ngrps; i++) {
            exp1[i] = MIN(exp[k], exp[k+1]);
            k += 2;
        }
    } else {
        for (i = 0, k = abs+1; i < ngrps; i++) {
            exp_min1 = MIN(exp[k  ], exp[k+1]);
            exp_min2 = MIN(exp[k+2], exp[k+3]);
            exp1[i]  = MIN(exp_min1, exp_min2);
            k += 4;
        }
    }

    // constraint for DC exponent
    if (!cpl)
        exp[abs] = MIN(exp[abs], 15);
    // Decrease the delta between each groups to within 2
    // so that they can be differentially encoded
    if (!cpl)
        exp1[0] = MIN(exp1[0], exp[0]+2);
    for (i = 1; i < ngrps; i++)
        exp1[i] = MIN(exp1[i], exp1[i-1]+2);
    for (i = ngrps-2; i >= 0; i--)
        exp1[i] = MIN(exp1[i], exp1[i+1]+2);

    if (cpl) {
        // coupling absolute exponent
        exp1[255] = exp1[0] - (exp1[0] & 1);
        exp[abs] = exp1[255];
    } else {
        exp[abs] = MIN(exp[abs], exp1[0]+2); // DC exponent is handled separately
    }

    if (grpsize == 2) {
        for (i = 0, k = abs+1; i < ngrps; i++) {
            v = exp1[i];
            exp[k] = v;
            exp[k+1] = v;
            k += 2;
        }
    } else {
        for (i = 0, k = abs+1; i < ngrps; i++) {
            v = exp1[i];
            exp[k] = v;
            exp[k+1] = v;
            exp[k+2] = v;
            exp[k+3] = v;
            k += 4;
        }
    }
}

/**
 * Determine a good exponent strategy for all blocks of a single channel.
 * A pre-defined set of strategies is chosen based on the SSE between each set
 * and the original exponents.
 */
static int compute_expstr_ch(uint8_t *exp[A52_NUM_BLOCKS], int start, int end,
                             int cpl, int search_size)
{
    int blk, s, str, i, j, k;
    int min_error, exp_error[A52_EXPSTR_SETS];
    int err;
    int ncoefs = end - start;
    uint8_t exponents[A52_NUM_BLOCKS][256];

    min_error = 0;
    for (s = 0; s < search_size; s++) {
        str = str_predef_priority[s];

        // collect exponents
        for (blk = 0; blk < A52_NUM_BLOCKS; blk++) {
            memcpy(exponents[blk], exp[blk], 256);
        }

        // encode exponents
        i = 0;
        while (i < A52_NUM_BLOCKS) {
            j = i + 1;
            while (j < A52_NUM_BLOCKS && str_predef[str][j]==EXP_REUSE) {
                exponent_min(exponents[i], exponents[j], ncoefs);
                j++;
            }
            encode_exp_blk_ch(exponents[i], start, end, str_predef[str][i], cpl);
            for (k = i+1; k < j; k++) {
                memcpy(exponents[k], exponents[i], 256);
            }
            i = j;
        }

        // select strategy based on minimum error from unencoded exponents
        exp_error[str] = 0;
        for (blk = 0; blk < A52_NUM_BLOCKS; blk++) {
            uint8_t *exp_blk = exp[blk];
            uint8_t *exponents_blk = exponents[blk];
            for (i = start; i < end; i++) {
                err = exp_blk[i] - exponents_blk[i];
                exp_error[str] += (err * err);
            }
        }
        if (exp_error[str] < exp_error[min_error])
            min_error = str;
    }
    return min_error;
}

/**
 * Runs the per-channel exponent strategy decision function for all channels
 */
static void compute_exponent_strategy(A52Context *ctx)
{
    A52Frame *frame = &ctx->frame;
    A52Block *blocks = frame->blocks;
    uint8_t *exp[A52_MAX_CHANNELS][A52_NUM_BLOCKS];
    int ch, blk, str;

    for (ch = !frame->cpl_in_use; ch <= ctx->n_channels; ch++) {
        str = str_predef_priority[0];
        if (ctx->params.expstr_search > 1) {
            for (blk = 0; blk < A52_NUM_BLOCKS; blk++)
                exp[ch][blk] = blocks[blk].exp[ch];
            str = compute_expstr_ch(exp[ch], frame->start_freq[ch],
                                    frame->end_freq[ch], (ch==CPL_CH),
                                    ctx->params.expstr_search);
        }
        for (blk = 0; blk < A52_NUM_BLOCKS; blk++)
            blocks[blk].exp_strategy[ch] = str_predef[str][blk];
        frame->expstr_set[ch] = str;
    }

    // lfe channel
    if (ctx->lfe) {
        for (blk = 0; blk < A52_NUM_BLOCKS; blk++)
            blocks[blk].exp_strategy[ctx->lfe_channel] = str_predef[0][blk];
    }
}

static void group_exp_blk_ch(uint8_t absexp, uint8_t *exp, int nexpgrps,
                             int gsize, uint8_t *grp_exp)
{
    int i;
    int delta[3];
    uint8_t exp0=0, exp1=0;
    uint8_t *p;

    p = exp;
    exp1 = absexp;
    grp_exp[0] = exp1;
                
    for (i = 1; i <= nexpgrps; i++) {
        // merge three delta into one code
        exp0 = exp1;
        exp1 = p[0];
        p += gsize;
        delta[0] = exp1 - exp0 + 2;

        exp0 = exp1;
        exp1 = p[0];
        p += gsize;
        delta[1] = exp1 - exp0 + 2;

        exp0 = exp1;
        exp1 = p[0];
        p += gsize;
        delta[2] = exp1 - exp0 + 2;

        grp_exp[i] = ((delta[0]*5+delta[1])*5)+delta[2];
    }
}

/**
 * Encode exponent groups.  3 exponents are in per 7-bit group.  The number of
 * groups varies depending on exponent strategy and bandwidth
 */
static void group_exponents(A52Context *ctx)
{
    A52Frame *frame = &ctx->frame;
    A52Block *block;
    int blk, ch;
    int expstr, gsize, bits, absexp;

    bits = 0;
    for (blk = 0; blk < A52_NUM_BLOCKS; blk++) {
        block = &frame->blocks[blk];
        for (ch = !frame->cpl_in_use; ch <= ctx->n_all_channels; ch++) {
            int cpl = (ch == CPL_CH);
            expstr = block->exp_strategy[ch];
            if (expstr == EXP_REUSE) {
                block->nexpgrps[ch] = 0;
                continue;
            }
            block->nexpgrps[ch] = nexpgrptab[cpl][expstr-1][frame->ncoefs[ch]];
            bits += (4 + (block->nexpgrps[ch] * 7));
            gsize = expstr + (expstr == EXP_D45);
            absexp = block->exp[ch][frame->start_freq[ch]-cpl];
            group_exp_blk_ch(absexp, &block->exp[ch][frame->start_freq[ch]+!cpl],
                             block->nexpgrps[ch], gsize, block->grp_exp[ch]);
        }
    }
    frame->exp_bits = bits;
}

/**
 * Creates final exponents for the entire frame based on exponent strategies.
 * If the strategy for a block & channel is EXP_REUSE, exponents are copied,
 * otherwise they are encoded according to the specific exponent strategy.
 */
static void encode_exponents(A52Context *ctx)
{
    A52Frame *frame = &ctx->frame;
    A52Block *blocks = frame->blocks;
    int ch, i, j, k;

    for (ch = !frame->cpl_in_use; ch <= ctx->n_all_channels; ch++) {
        // compute the exponents as the decoder will see them. The
        // EXP_REUSE case must be handled carefully : we select the
        // min of the exponents
        i = 0;
        while (i < A52_NUM_BLOCKS) {
            j = i + 1;
            while (j < A52_NUM_BLOCKS && blocks[j].exp_strategy[ch]==EXP_REUSE) {
                exponent_min(&blocks[i].exp[ch][frame->start_freq[ch]],
                             &blocks[j].exp[ch][frame->start_freq[ch]],
                             frame->ncoefs[ch]);
                j++;
            }
            encode_exp_blk_ch(blocks[i].exp[ch], frame->start_freq[ch],
                              frame->end_freq[ch],
                              blocks[i].exp_strategy[ch], (ch==CPL_CH));
            // copy encoded exponents for reuse case
            for (k = i+1; k < j; k++) {
                memcpy(&blocks[k].exp[ch][frame->start_freq[ch]],
                       &blocks[i].exp[ch][frame->start_freq[ch]],
                       frame->ncoefs[ch]);
            }
            i = j;
        }
    }
}

/**
 * Extracts the optimal exponent portion of each MDCT coefficient.
 */
static void extract_exponents(A52Context *ctx)
{
    A52Frame *frame = &ctx->frame;
    A52Block *block;
    int all_channels = ctx->n_all_channels;
    int blk, ch, j;
    uint32_t v;

    for (ch = !frame->cpl_in_use; ch <= all_channels; ch++) {
        for (blk = 0; blk < A52_NUM_BLOCKS; blk++) {
            block = &frame->blocks[blk];
            for (j = frame->start_freq[ch]; j < 256; j++) {
                v = (uint32_t)fabs(block->mdct_coef[ch][j] * 16777216.0);
                block->exp[ch][j] = (v == 0)? 24 : 23 - log2i(v);
            }
        }
    }
}

/**
 * Runs all the processes in extracting, analyzing, and encoding exponents
 */
void a52_exponent_process(A52Context *ctx)
{
    extract_exponents(ctx);

    compute_exponent_strategy(ctx);

    encode_exponents(ctx);

    group_exponents(ctx);
}
